package com.lbw.controller.chat;

import com.lbw.chatMemory.repository.MySqlRAGMemoryRepository;
import com.lbw.entity.vo.Result;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Objects;

@Slf4j
@RequiredArgsConstructor
@RestController
@RequestMapping("/pdf")
public class RAGController {

    // 本地部署的大模型
    @Autowired
    @Qualifier("pdfChatClient")
    private ChatClient pdfChatClient;

    // 调用远程大模型api
    @Autowired
    private ChatClient openAiPdfChatClient;

    // mysql存储
    @Autowired
    private MySqlRAGMemoryRepository mySqlRAGMemoryRepository;

    // redis向量库，ollama本地部署的向量化模型
    @Autowired
    private RedisVectorStore redisVectorStore;

    @GetMapping(value = "/chat", produces = "text/html;charset=UTF-8")
    public Flux<String> chat(String prompt, String chatId) {
        // 保存会话id
        mySqlRAGMemoryRepository.saveConversationId(chatId);

        // 调用本地模型
        return pdfChatClient
                .prompt(prompt)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                //.advisors(a -> a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION, "file_name == '"+file.getFilename()+"'"))
                .stream()
                .content();

        // 调用远程模型
        /*return openAiPdfChatClient
                .prompt(prompt)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                .stream()
                .content();*/
    }

    /**
     * 文件上传
     */
    @RequestMapping("/upload/{chatId}")
    public Result uploadPdf(@PathVariable String chatId, @RequestParam("file") MultipartFile file) {
        try {
            // 校验文件是否为PDF格式
            if (Objects.equals(file.getContentType(), "application/pdf")) {
                // 1.创建PDF的读取器
                PagePdfDocumentReader reader = new PagePdfDocumentReader(
                        file.getResource(), // 文件源
                        PdfDocumentReaderConfig.builder()
                                .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                                .withPagesPerDocument(1) // 每1页PDF作为一个Document
                                .build()
                );
                // 2.读取PDF文档，拆分为Document
                List<Document> documents = reader.read();
                // 3.写入向量库
                redisVectorStore.add(documents); // redis向量数据库，调用ollama的向量模型
            }else {
                // 写入向量库
                TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(file.getResource());
                List<Document> documents = tikaDocumentReader.read();
                // 分割文本
                TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
                List<Document> apply = tokenTextSplitter.apply(documents);
                redisVectorStore.add(apply);
            }
            return Result.ok();
        } catch (Exception e) {
            log.error("Failed to upload PDF.", e);
            return Result.fail("上传文件失败！");
        }
    }

}
